import csv
import queue


def get_skeleton_data():
    # 初始化Pygame
    pygame.init()

    # 设置窗口尺寸
    width, height = 960, 540
    size = (width, height)
    screen = pygame.display.set_mode(size)
    pygame.display.set_caption("PyKinect2 Body Tracking")  # 窗口标题

    # 创建Kinect实例
    # kinect = PyKinectRuntime.PyKinectRuntime(PyKinectV2.FrameSourceTypes_Body)
    kinect = PyKinectRuntime.PyKinectRuntime(PyKinectV2.FrameSourceTypes_Color | PyKinectV2.FrameSourceTypes_Body)

    # 创建CSV文件
    with open('skeleton_data.csv', 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        joint_ids = [joint_id for joint_id in range(0, PyKinectV2.JointType_Count)]
        header = ['Frame'] + \
                 [f'Joint {joint_id + 1} X' for joint_id in joint_ids] + \
                 [f'Joint {joint_id + 1} Y' for joint_id in joint_ids] + \
                 [f'Joint {joint_id + 1} Z' for joint_id in joint_ids]
        writer.writerow(header)  # 写入表头

        # 游戏主循环
        running = True
        frame_count = 0
        frame_rate = 30  # 帧率设定为30帧/秒
        interval = 1.0 / frame_rate  # 计算每帧的时间间隔
        last_time = pygame.time.get_ticks() / 1000.0  # 上一帧的时间戳
        tracked_body_index = None  # 跟踪的人体索引

        while running:
            # 处理事件
            # for event in pygame.event.get():
            #     if event.type == pygame.QUIT:
            #         running = False

            # 获取骨架数据
            if kinect.has_new_body_frame():
                bodies = kinect.get_last_body_frame().bodies
                if bodies is not None:
                    # 寻找正在跟踪的人体
                    for i in range(0, kinect.max_body_count):
                        body = bodies[i]
                        if body.is_tracked:
                            tracked_body_index = i
                            break

            # 跟踪到人体后才进行数据记录
            if tracked_body_index is not None:
                # 计算当前时间戳
                current_time = pygame.time.get_ticks() / 1000.0

                # 根据时间间隔记录数据
                if current_time - last_time >= interval:
                    body = bodies[tracked_body_index]
                    frame_data = []

                    # 存储关节坐标数据
                    # 按照时序将25个骨架识别点的坐标，以“Joint1X、Joint2X、
                    # ...、Joint25X、Joint1Y、Joint2Y.....Joint25Z”的顺序，
                    # 按行逐帧存储在csv文件中
                    for joint_id in range(0, PyKinectV2.JointType_Count):  # JointType_Count = 25
                        joint = body.joints[joint_id]
                        frame_data.extend([joint.Position.x, joint.Position.y, joint.Position.z])

                    # 调整关节坐标数据的排列顺序
                    rearranged_data = []
                    for joint_id in range(1, PyKinectV2.JointType_Count + 1):
                        rearranged_data.extend(frame_data[(joint_id - 1) * 3: joint_id * 3])

                    # 写入CSV文件
                    writer.writerow([frame_count + 1] + rearranged_data)

                    # 打印关节数据
                    print(f"Frame: {frame_count + 1}, Joint Positions: {rearranged_data}")

                    # 更新上一帧的时间戳
                    last_time = current_time

                    # 帧数加1
                    frame_count += 1

            # 清空屏幕
            screen.fill((0, 0, 0))

            # 绘制骨架连线图
            if tracked_body_index is not None:
                joints = bodies[tracked_body_index].joints

                # 定义骨架连线关系
                bone_connections = [
                    # (JointType_Head, JointType_SpineShoulder),
                    # (JointType_SpineShoulder, JointType_SpineMid),
                    # (JointType_SpineMid, JointType_SpineBase),
                    # (JointType_SpineShoulder, JointType_ShoulderRight),
                    # (JointType_SpineShoulder, JointType_ShoulderLeft),
                    # (JointType_ShoulderRight, JointType_ElbowRight),
                    # (JointType_ElbowRight, JointType_WristRight),
                    # (JointType_ShoulderLeft, JointType_ElbowLeft),
                    # (JointType_ElbowLeft, JointType_WristLeft)

                    (JointType_Head, JointType_Neck),
                    (JointType_Neck, JointType_SpineShoulder),
                    (JointType_SpineShoulder, JointType_SpineMid),
                    (JointType_SpineMid, JointType_SpineBase),
                    (JointType_SpineShoulder, JointType_ShoulderRight),
                    (JointType_SpineShoulder, JointType_ShoulderLeft),
                    (JointType_ShoulderRight, JointType_ElbowRight),
                    (JointType_ElbowRight, JointType_WristRight),
                    (JointType_WristRight, JointType_HandRight),
                    (JointType_ShoulderLeft, JointType_ElbowLeft),
                    (JointType_ElbowLeft, JointType_WristLeft),
                    (JointType_WristLeft, JointType_HandLeft),
                    (JointType_SpineBase, JointType_HipRight),
                    (JointType_SpineBase, JointType_HipLeft),
                    (JointType_HipRight, JointType_KneeRight),
                    (JointType_KneeRight, JointType_AnkleRight),
                    (JointType_AnkleRight, JointType_FootRight),
                    (JointType_HipLeft, JointType_KneeLeft),
                    (JointType_KneeLeft, JointType_AnkleLeft),
                    (JointType_AnkleLeft, JointType_FootLeft)
                ]

                # 绘制骨架连线
                for connection in bone_connections:
                    joint1 = joints[connection[0]]
                    joint2 = joints[connection[1]]
                    if joint1.TrackingState == TrackingState_Tracked and joint2.TrackingState == TrackingState_Tracked:
                        joint1_pos = kinect.body_joint_to_color_space(joint1)
                        joint2_pos = kinect.body_joint_to_color_space(joint2)
                        pygame.draw.line(screen, (255, 0, 0), (joint1_pos.x, joint1_pos.y),
                                         (joint2_pos.x, joint2_pos.y), 5)

            # 更新屏幕显示
            pygame.display.flip()

        # 退出Pygame
        pygame.quit()

    # 释放Kinect资源
    kinect.close()


from pykinect2 import PyKinectV2
from pykinect2.PyKinectV2 import *
from pykinect2 import PyKinectRuntime

import ctypes

import pygame
import sys
from multiprocessing import Queue

if sys.hexversion >= 0x03000000:
    pass
else:
    pass

# colors for drawing different bodies
SKELETON_COLORS = [pygame.color.THECOLORS["red"],
                   pygame.color.THECOLORS["blue"],
                   pygame.color.THECOLORS["green"],
                   pygame.color.THECOLORS["orange"],
                   pygame.color.THECOLORS["purple"],
                   pygame.color.THECOLORS["yellow"],
                   pygame.color.THECOLORS["violet"]]


class BodyGameRuntime(object):
    def __init__(self):
        pygame.init()

        # Used to manage how fast the screen updates
        self._clock = pygame.time.Clock()

        # Set the width and height of the screen [width, height]
        self._infoObject = pygame.display.Info()
        self._screen = pygame.display.set_mode((self._infoObject.current_w >> 1, self._infoObject.current_h >> 1),
                                               pygame.HWSURFACE | pygame.DOUBLEBUF | pygame.RESIZABLE, 32)

        pygame.display.set_caption("Kinect for Windows v2")

        # Loop until the user clicks the close button.
        self._done = False

        # Used to manage how fast the screen updates
        self._clock = pygame.time.Clock()

        # Kinect runtime object, we want only color and body frames
        self._kinect = PyKinectRuntime.PyKinectRuntime(
            PyKinectV2.FrameSourceTypes_Color | PyKinectV2.FrameSourceTypes_Body)

        # back buffer surface for getting Kinect color frames, 32bit color, width and height equal to the Kinect color frame size
        self._frame_surface = pygame.Surface(
            (self._kinect.color_frame_desc.Width, self._kinect.color_frame_desc.Height), 0, 32)

        # here we will store skeleton data
        self._bodies = None

    def draw_body_bone(self, joints, jointPoints, color, joint0, joint1):
        joint0State = joints[joint0].TrackingState
        joint1State = joints[joint1].TrackingState

        # both joints are not tracked
        if (joint0State == PyKinectV2.TrackingState_NotTracked) or (
                joint1State == PyKinectV2.TrackingState_NotTracked):
            return

        # both joints are not *really* tracked
        if (joint0State == PyKinectV2.TrackingState_Inferred) and (
                joint1State == PyKinectV2.TrackingState_Inferred):
            return

        # ok, at least one is good
        start = (jointPoints[joint0].x, jointPoints[joint0].y)
        end = (jointPoints[joint1].x, jointPoints[joint1].y)

        try:
            pygame.draw.line(self._frame_surface, color, start, end, 8)
        except:  # need to catch it due to possible invalid positions (with inf)
            pass

    def draw_body(self, joints, jointPoints, color):
        """
        画出人体骨骼线
        """
        # Torso 躯干
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_Head, PyKinectV2.JointType_Neck)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_Neck,
                            PyKinectV2.JointType_SpineShoulder)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_SpineShoulder,
                            PyKinectV2.JointType_SpineMid)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_SpineMid,
                            PyKinectV2.JointType_SpineBase)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_SpineShoulder,
                            PyKinectV2.JointType_ShoulderRight)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_SpineShoulder,
                            PyKinectV2.JointType_ShoulderLeft)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_SpineBase,
                            PyKinectV2.JointType_HipRight)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_SpineBase,
                            PyKinectV2.JointType_HipLeft)

        # Right Arm
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_ShoulderRight,
                            PyKinectV2.JointType_ElbowRight)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_ElbowRight,
                            PyKinectV2.JointType_WristRight)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_WristRight,
                            PyKinectV2.JointType_HandRight)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_HandRight,
                            PyKinectV2.JointType_HandTipRight)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_WristRight,
                            PyKinectV2.JointType_ThumbRight)

        # Left Arm
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_ShoulderLeft,
                            PyKinectV2.JointType_ElbowLeft)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_ElbowLeft,
                            PyKinectV2.JointType_WristLeft)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_WristLeft,
                            PyKinectV2.JointType_HandLeft)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_HandLeft,
                            PyKinectV2.JointType_HandTipLeft)
        self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_WristLeft,
                            PyKinectV2.JointType_ThumbLeft)

        # Right Leg
        # self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_HipRight,
        #                     PyKinectV2.JointType_KneeRight)
        # self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_KneeRight,
        #                     PyKinectV2.JointType_AnkleRight)
        # self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_AnkleRight,
        #                     PyKinectV2.JointType_FootRight)

        # Left Leg
        # self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_HipLeft,
        #                     PyKinectV2.JointType_KneeLeft)
        # self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_KneeLeft,
        #                     PyKinectV2.JointType_AnkleLeft)
        # self.draw_body_bone(joints, jointPoints, color, PyKinectV2.JointType_AnkleLeft,
        #                     PyKinectV2.JointType_FootLeft)

    def draw_color_frame(self, frame, target_surface):
        target_surface.lock()
        address = self._kinect.surface_as_array(target_surface.get_buffer())
        ctypes.memmove(address, frame.ctypes.data, frame.size)
        del address
        target_surface.unlock()

    # def run(self,queue:Queue):
    def run(self, kinect_waypoint: Queue):
        # -------- Main Program Loop -----------

        while not self._done:
            # --- Main event loop
            for event in pygame.event.get():  # User did something
                if event.type == pygame.QUIT:  # If user clicked close
                    self._done = True  # Flag that we are done so we exit this loop

                elif event.type == pygame.VIDEORESIZE:  # window resized
                    self._screen = pygame.display.set_mode(event.dict['size'],
                                                           pygame.HWSURFACE | pygame.DOUBLEBUF | pygame.RESIZABLE,
                                                           32)

            # --- Game logic should go here

            # --- Getting frames and drawing
            # --- Woohoo! We've got a color frame! Let's fill out back buffer surface with frame's data
            if self._kinect.has_new_color_frame():
                frame = self._kinect.get_last_color_frame()
                self.draw_color_frame(frame, self._frame_surface)
                frame = None

            # --- Cool! We have a body frame, so can get skeletons
            if self._kinect.has_new_body_frame():
                self._bodies = self._kinect.get_last_body_frame()

            # --- draw skeletons to _frame_surface
            frame_data = []
            # kinect_waypoint.get()
            if self._bodies is not None:
                for i in range(0, self._kinect.max_body_count):
                    body = self._bodies.bodies[i]
                    if not body.is_tracked:
                        continue
                    # 存储关节坐标数据
                    # 按照时序将25个骨架识别点的坐标，以“Joint1X、Joint2X、
                    # ...、Joint25X、Joint1Y、Joint2Y.....Joint25Z”的顺序，
                    # frame_data = []
                    for joint_id in range(0, PyKinectV2.JointType_Count):  # JointType_Count = 25
                        joint = body.joints[joint_id]
                        if joint_id == 0 or joint_id == 1 or joint_id == 3 or joint_id == 4 or joint_id == 5 or joint_id == 6 or joint_id == 8 or joint_id == 9 or joint_id == 10 or joint_id == 20:
                            # extend()是列表对象的一个方法，用于将一个可迭代对象（如列表、元组等）的元素添加到列表的末尾。
                            # frame_data.extend([joint.Position.x, joint.Position.y, joint.Position.z])
                            frame_data.extend([joint.Position.z - 1.3, joint.Position.x + 0.5, joint.Position.y + 0.2])   # 对齐机械臂基座坐标
                    # kinect_waypoint.put(frame_data) # 向管道存放当前骨骼坐标信息，列表
                    # # print(f"Joint Positions:{frame_data}")
                    # # time.sleep(1)

                    # convert joint coordinates to color space
                    joints = body.joints
                    joint_points = self._kinect.body_joints_to_color_space(joints)
                    self.draw_body(joints, joint_points, SKELETON_COLORS[i])
            else:
                frame_data = [0]

            # 将此刻kinect_waypoint管道（骨骼坐标）清空，并put最新骨骼坐标
            while kinect_waypoint.empty() is not True:
                try:
                    kinect_waypoint.get(timeout=0.1)
                except queue.Empty:
                    pass

            if kinect_waypoint.empty() is True:
                try:
                    kinect_waypoint.put(frame_data,timeout=0.1)  # 向管道存放当前骨骼坐标信息，列表
                except queue.Full:
                    pass
                    # print("TimeOut: kinect_waypoint_queue is already fully")


            # print(f"Joint Positions:{frame_data}")
            # time.sleep(1)
            # --- copy back buffer surface pixels to the screen, resize it if needed and keep aspect ratio
            # --- (screen size may be different from Kinect's color frame size)
            h_to_w = float(self._frame_surface.get_height()) / self._frame_surface.get_width()
            target_height = int(h_to_w * self._screen.get_width())
            surface_to_draw = pygame.transform.scale(self._frame_surface,
                                                     (self._screen.get_width(), target_height));
            self._screen.blit(surface_to_draw, (0, 0))
            surface_to_draw = None
            pygame.display.update()

            # --- Go ahead and update the screen with what we've drawn.
            pygame.display.flip()

            # --- Limit to 60 frames per second
            self._clock.tick(60)

            # for joint_id in range(0, PyKinectV2.JointType_Count):
            #     joint = body.joints[joint_id]
            #     print(joint)

        # Close our Kinect sensor, close the window and quit.
        self._kinect.close()
        pygame.quit()


if __name__ == "__main__":
    # get_skeleton_data() # 这是两个程序
    q = Queue()
    game = BodyGameRuntime()
    game.run(q)
